from MSTH.configs.method_import import *
import numpy as np
import itertools
import random

method_configs: Dict[str, Union[TrainerConfig, VideoTrainerConfig]] = {}

method_configs["base_it20000_base_32"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=10,
            static_ratio_decay_total_steps=500000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=32,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=5e-4, max_steps=40000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=5e-4, max_steps=40000),
        },
    },
)

method_configs["base_it20000_nodownsample"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos"),
                # data=Path("/data/machine/data/fit/videos_2"),
                # downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 2 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=10,
            static_ratio_decay_total_steps=500000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=10,
            static_ratio_decay_total_steps=500000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform2"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="stratified",
            # n_time_for_dynamic=3,
            # n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=500000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=500000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform4"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 10 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=500000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform5"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 3 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=500000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform6"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 7 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=500000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform7"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="stratified",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 7 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=10,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform8"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="stratified",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 7 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=40,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)


method_configs["base_it20000_uniform9"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="stratified",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 7 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=30,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform10"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="stratified",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 7 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=20,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform11"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="stratified",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 7 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=5,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform12"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 3 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=40,
            static_dynamic_sampling_ratio_end=40,
            static_ratio_decay_total_steps=500000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform13"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 3 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=60,
            static_dynamic_sampling_ratio_end=60,
            static_ratio_decay_total_steps=500000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3_16"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=500000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=16,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)


method_configs["base_it20000_uniform3_sharpen2"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_salmon_videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=500000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            # dist_sharpen=1.3,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3_cook_spinach"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=500000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)


method_configs["base_it20000_uniform3_cook_spinach_50_10"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=10,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3_cook_spinach_10_2"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=10,
            static_dynamic_sampling_ratio_end=2,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.01,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3_coffee_martini_10_2"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=10,
            static_dynamic_sampling_ratio_end=2,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.01,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3_cook_spinach_5_1"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=5,
            static_dynamic_sampling_ratio_end=1,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.01,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3_cook_spinach_5_1_20"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 20 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=5,
            static_dynamic_sampling_ratio_end=1,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.01,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3_cook_spinach_5_1_bighash"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=5,
            static_dynamic_sampling_ratio_end=1,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=22,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3_cut_roasted_beef"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cut_roasted_beef/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)


method_configs["base_it20000_uniform3_flame_steak"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/flame_steak/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=10,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)


method_configs["base_it20000_uniform3_coffee_martini"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=10,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3_sear_steak"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/sear_steak/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=10,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 15),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1.0,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)


base_method = method_configs["base_it20000_uniform3_cook_spinach_10_2"]


def setp(exps):
    def setfunc(x, v):
        if isinstance(exps, (tuple, list)):
            for exp in exps:
                command = "x." + exp + "=v"
                exec(command)
        else:
            command = "x." + exps + "=v"
            exec(command)

    return setfunc


set_functions = {
    "n_time": setp("pipeline.datamanager.n_time_for_dynamic"),
    "sampling_ratio_start": setp("pipeline.datamanager.static_dynamic_sampling_ratio"),
    "sampling_ratio_end": setp("pipeline.datamanager.static_dynamic_sampling_ratio_end"),
    "sampling_ratio_decay": setp("pipeline.datamanager.static_ratio_decay_total_steps"),
    "mask_loss_mult": setp("pipeline.model.mask_loss_mult"),
    "mask_loss_for_proposal": setp("pipeline.model.mask_loss_for_proposal"),
    "mask_init_mean": setp(
        ["pipeline.model.mask_init_mean", "pipeline.model.proposal_net_args_list[0]['mask_init_mean']"]
    ),
    "base_reso": setp("pipeline.model.base_res"),
}

potential_values = {
    # "n_time": [
    # lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (40000 - 1000))),
    # ],
    # "sampling_ratio_start": [50, 40, 30, 20, 10],
    # "sampling_ratio_end": [40, 30, 20, 10, 1],
    "n_time": [
        lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
        lambda x: 1 if x < 1000 else 1 + 10 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
        lambda x: 1 if x < 1000 else 1 + 15 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
    ],
    "mask_loss_mult": [
        0.01,
        0.1,
    ],
    "mask_loss_for_proposal": [
        True,
        False,
    ],
    "base_reso": [
        (16, 16, 16, 15),
        (16, 16, 16, 30),
        (16, 16, 16, 60),
        (16, 16, 16, 120),
    ],
}

all_hyper_parameter_key = potential_values.keys()
all_hyper_parameter_value = [potential_values[k] for k in all_hyper_parameter_key]
all_specs = list(itertools.product(*all_hyper_parameter_value))
all_specs = [{k: v for k, v in zip(all_hyper_parameter_key, spec)} for spec in all_specs]
random.shuffle(all_specs)
print("==== ALL SPECS ====")
print(all_specs)

for i, spec in enumerate(all_specs):
    method_configs[f"anoynmous_method_{i}"] = copy.deepcopy(base_method)
    for k, v in spec.items():
        set_functions[k](method_configs[f"anoynmous_method_{i}"], v)
    print(method_configs[f"anoynmous_method_{i}"])
# set_sampling_ratio_decay(anoynmous_method, anoynmous_method.max_num_iterations)


method_configs["base_it20000_uniform3_cook_spinach_high_dynamic_base"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=10,
            static_dynamic_sampling_ratio_end=2,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3_cook_spinach_high_dynamic_base_longer"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=40000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=10,
            static_dynamic_sampling_ratio_end=2,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
    },
)


method_configs["base_it20000_uniform3_cook_spinach_high_dynamic_base_longer2"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=40000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=10,
            static_dynamic_sampling_ratio_end=1,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
    },
)

method_configs["base_it40000_uniform3_cook_spinach_high_dynamic_base_new_50_30"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=40000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=30,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
    },
)

method_configs["base_it40000_uniform3_cook_spinach_high_dynamic_base_new_50_30_mst"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=40000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=30,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
    },
)

method_configs[
    "base_it40000_uniform3_cook_spinach_high_dynamic_base_new_50_10_mst_independent"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=30,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="independent",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "independent",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it40000_uniform3_coffee_martini_high_dynamic_base_new_50_30_mst"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=40000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=30,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
    },
)

method_configs["base_it20000_uniform3_coffee_martini_high_dynamic_base_new_50_30_mst2"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=30,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3_coffee_martini_high_dynamic_base_new_50_30_mst3"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=70,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs[
    "base_it20000_uniform3_coffee_martini_high_dynamic_base_new_50_30_mst3_sc1"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.25,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            # n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=70,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3_coffee_martini_sim1"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=10000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            # n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (10000 - 1000))),
            static_dynamic_sampling_ratio=70,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=10000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=15000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=15000),
        },
    },
)

method_configs["base_it20000_uniform3_coffee_martini_sim2"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=10000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            # n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (10000 - 1000))),
            static_dynamic_sampling_ratio=70,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=10000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(1024,),
            num_nerf_samples_per_ray=128,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=15000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=15000),
        },
    },
)


method_configs[
    "base_it20000_uniform3_coffee_martini_high_dynamic_base_new_50_30_mst3_l2contraction"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=70,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            contraction_type="l2",
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs[
    "base_it20000_uniform3_coffee_martini_high_dynamic_base_new_50_30_mst35"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=70,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            middle_distance=2.0,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)


method_configs[
    "base_it20000_uniform3_coffee_martini_high_dynamic_base_new_50_30_mst36"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=70,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            middle_distance=0.7,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs[
    "base_it20000_uniform3_coffee_martini_high_dynamic_base_new_50_30_mst36_nocontract"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
                scene_scale=8,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=70,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            disable_scene_contraction=True,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            middle_distance=0.7,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs[
    "base_it20000_uniform3_coffee_martini_high_dynamic_base_new_50_30_mst37"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=70,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 120),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            middle_distance=0.7,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs[
    "base_it20000_uniform3_coffee_martini_high_dynamic_base_new_50_30_mst34"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=70,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(1024, 1024, 384, 100),
            base_res=(64, 64, 32, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(64, 64, 64),
            mask_log2_hash_size=18,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            hidden_dim=128,
            hidden_dim_color=128,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 17,
                    "log2_hashmap_size_temporal": 17,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs[
    "base_it20000_uniform3_coffee_martini_high_dynamic_base_new_50_30_mst33"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=70,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(1024, 1024, 1024, 100),
            base_res=(32, 32, 32, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=19,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(64, 64, 64),
            mask_log2_hash_size=18,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            hidden_dim=128,
            hidden_dim_color=128,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 17,
                    "log2_hashmap_size_temporal": 17,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs[
    "base_it20000_uniform3_coffee_martini_high_dynamic_base_new_50_30_mst32"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=70,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            hidden_dim=128,
            hidden_dim_color=128,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs[
    "base_it20000_uniform3_coffee_martini_high_dynamic_base_new_50_30_mst31"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 1 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=70,
            static_dynamic_sampling_ratio_end=50,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3_coffee_martini_high_dynamic_base_new_50_30_mst4"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/coffee_martini/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=30,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(256,),
            num_nerf_samples_per_ray=64,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=2e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs[
    "base_it40000_uniform3_cook_spinach_high_dynamic_base_new_50_30_mst_useweightdiff"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=40000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (40000 - 1000))),
            use_temporal_weight="diff",
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=30,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
    },
)

method_configs[
    "base_it40000_uniform3_cook_spinach_high_dynamic_base_new_50_30_mst_useweight"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=40000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (40000 - 1000))),
            use_temporal_weight="median",
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=30,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mst",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mst",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
    },
)

method_configs["base_it20000_uniform3_cook_spinach_high_dynamic_base_longer4"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=40000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=10,
            static_dynamic_sampling_ratio_end=1,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=21,
            log2_hashmap_size_temporal=23,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
    },
)

method_configs["base_it20000_uniform3_cook_spinach_high_dynamic_base_longer5"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=40000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=10,
            static_dynamic_sampling_ratio_end=1,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=20,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
    },
)

method_configs["base_it20000_uniform3_cook_spinach_high_dynamic_base_hierarchical"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=40000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=10,
            static_dynamic_sampling_ratio_end=3,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="hierarchical",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 17,
                    "log2_hashmap_size_temporal": 17,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
    },
)

method_configs["base_it20000_uniform3_cook_spinach_high_dynamic_base_hierarchical2"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=40000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=30,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="hierarchical",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 17,
                    "log2_hashmap_size_temporal": 17,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
    },
)

method_configs[
    "base_it20000_uniform3_cook_spinach_high_dynamic_base_hierarchical2_20000"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=30,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="hierarchical",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 17,
                    "log2_hashmap_size_temporal": 17,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs[
    "base_it20000_uniform3_cook_spinach_high_dynamic_base_hierarchical2_20000_mask_loss"
] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=30,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=1,
            use_no_prior_mask=False,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="hierarchical",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 17,
                    "log2_hashmap_size_temporal": 17,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)

method_configs["base_it20000_uniform3_cook_spinach_high_dynamic_base_hierarchical3"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=40000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=50,
            static_dynamic_sampling_ratio_end=30,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="hierarchical",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 17,
                    "log2_hashmap_size_temporal": 17,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
    },
)

method_configs["base_it20000_uniform3_cook_spinach_high_dynamic_base_longer3"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=40000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=10,
            static_dynamic_sampling_ratio_end=1,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 300),
                    "base_res": (16, 16, 16, 30),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=75000),
        },
    },
)

method_configs["base_it20000_uniform3_cook_spinach_high_dynamic_base_prior"] = SpaceTimeHashingTrainerConfig(
    method_name="Spatial_Time_Hashing_With_Base",
    steps_per_eval_batch=1000,
    steps_per_save=50000,
    max_num_iterations=20000,
    mixed_precision=True,
    log_gradients=True,
    pipeline=SpaceTimePipelineConfig(
        datamanager=SpaceTimeDataManagerConfig(
            dataparser=VideoDataParserConfig(
                data=Path("/data/machine/data/cook_spinach/videos_2/"),
                # data=Path("/data/machine/data/fit/videos_2"),
                downscale_factor=2,
                # scale_factor=1.0 / 2.0,
                scale_factor=0.5,
            ),
            train_num_rays_per_batch=16384,
            # eval_num_rays_per_batch=32768,
            camera_optimizer=CameraOptimizerConfig(mode="off"),
            use_uint8=True,
            # use_stratified_pixel_sampler=True,
            spatial_temporal_sampler="st",
            # n_time_for_dynamic=3,
            n_time_for_dynamic=lambda x: 1 if x < 1000 else 1 + 5 * np.sin((x - 1000) * np.pi / (2 * (20000 - 1000))),
            static_dynamic_sampling_ratio=10,
            static_dynamic_sampling_ratio_end=2,
            static_ratio_decay_total_steps=20000,
        ),
        model=DSpaceTimeHashingModelConfig(
            # distortion_loss_mult=0.0,
            max_res=(2048, 2048, 2048, 300),
            base_res=(16, 16, 16, 30),
            # num_proposal_samples_per_ray=(256, 96),
            num_proposal_samples_per_ray=(128,),
            num_nerf_samples_per_ray=48,
            proposal_weights_anneal_max_num_iters=5000,
            # proposal_weights_anneal_slope = 10.0,
            log2_hashmap_size_spatial=19,
            log2_hashmap_size_temporal=21,
            eval_num_rays_per_chunk=32768,
            mask_loss_mult=0.1,
            mask_loss_for_proposal=False,
            mst_mode="mt",
            mask_reso=(128, 128, 128),
            mask_log2_hash_size=21,
            mask_type="global",
            st_mlp_mode="shared",
            num_proposal_iterations=1,
            use_loss_static=False,
            render_static=False,
            interp="linear",
            mask_init_mean=-1,
            use_no_prior_mask=False,
            proposal_net_args_list=[
                {
                    "hidden_dim": 16,
                    "log2_hashmap_size_spatial": 18,
                    "log2_hashmap_size_temporal": 18,
                    "num_levels": 5,
                    "max_res": (128, 128, 128, 150),
                    "base_res": (16, 16, 16, 15),
                    "use_linear": False,
                    "mode": "mt",
                    "mask_reso": (64, 64, 64),
                    "mask_log2_hash_size": 18,
                    "mask_type": "global",
                    "st_mlp_mode": "shared",
                    "interp": "linear",
                    "mask_init_mean": -1,
                },
            ],
        ),
    ),
    optimizers={
        "proposal_networks": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
        "fields": {
            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
            "scheduler": ExponentialDecaySchedulerConfig(lr_final=2e-4, max_steps=35000),
        },
    },
)
